import json
from fastapi import APIRouter
from starlette.websockets import WebSocket
from client.openai_client import invoke_openai, invoke_laozhang_api
from dao.es_dao.p_health_log_dao import find_health_log_dao

ai_router = APIRouter(prefix="/test/ai", tags=["人工智能接口"])


@ai_router.websocket("/healthAnalysisWs")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    try:
        while True:
            # 接收前端参数
            raw_data = await ws.receive_text()
            data = json.loads(raw_data)
            health_data = await find_health_log_dao(data)
            query = f"请帮我分析一下这些健康数据，是同一个人的，请根据时间轴做出科学的判断。要求简单且高大上，用户是保洁公司的领导，你直接面对的是用户，不要说多余的话，因为是展示在页面上的内容" \
                    f"健康数据是保洁员工的。年龄大多在60岁以上，因此要根据实际情况进行科学分析，答案以'以下是该名员工在{data['startDate']}到{data['endDate']}之间的健康数据分析，" \
                    f"要求结果简单明了，不要加入其他的思考内容。再给出一些健康建议，和注意事项，作为保洁公司领导，要如何做出一些决策，要充分考虑用人风险和关怀。以下是数据JSON格式。"
            data_query = f"{query}{json.dumps(health_data)}"
            model = 'gpt-4o-mini'
            if data.get('model'):
                model = data.get('model')
            await invoke_laozhang_api(data_query, ws, model=model)
    except Exception as e:
        await ws.close()
        print(f"WebSocket 错误: {str(e)}")
